#ifndef _ELLGRAPH_HXX__
#define _ELLGRAPH_HXX__

// #define DEBUG
#include "config.hxx"
#include "../utils.hxx"
#include "GraphData.cuh"

namespace GraphData {

template <typename IndexType>
struct ELLGraph: public MetaGraph<IndexType> {
  IndexType *col_idx;
  IndexType *row_idx;

  IndexType *col_eidx;
  IndexType *row_eidx;
  
  IndexType col_stride;
  IndexType row_stride;

  using meta = MetaGraph<IndexType>;

  ELLGraph( size_t nvertex, size_t nedge, DeviceTag on_dev,
            IndexType col_stride, IndexType *col_idx, IndexType *col_eidx,
            IndexType row_stride = 0, IndexType *row_idx = nullptr, IndexType *row_eidx = nullptr) 
            : MetaGraph<IndexType>(nvertex, nedge, on_dev), col_stride(col_stride), col_idx(col_idx),
              col_eidx(col_eidx), row_stride(row_stride), row_idx(row_idx), row_eidx(row_eidx) {
    
    meta::independent = false;

    assert(col_idx != nullptr && col_stride > 0);
    if ( row_stride ) {
      assert(row_idx != nullptr);
      assert(col_eidx != nullptr && row_eidx != nullptr 
             || col_eidx != nullptr && row_eidx != nullptr);
    }
  }

  ELLGraph( IndexType *src_raw, IndexType *dst_raw, IndexType *eidx_raw, 
            size_t nedge, bool with_reverse = false)
          : MetaGraph<IndexType>(src_raw, dst_raw, eidx_raw, nedge) {
    if ( with_reverse ) row_idx = reinterpret_cast<IndexType*>(100);
    fromCOO(src_raw, dst_raw, eidx_raw);
  }

  void fromCOO(IndexType src_raw, IndexType dst_raw, IndexType *eidx_raw ) {
    IndexType *odeg, *ideg, *ideg_, *odeg_;

    bool with_reverse = (row_idx != nullptr);
    bool with_weight = meta::weighted;
    row_stride = col_stride = -1;

    MEMALLOC(IndexType, ideg, meta::nvertex, meta::on_dev);
    MEMALLOC(IndexType, odeg, meta::nvertex, meta::on_dev);
    MEMCLEAN(ideg, meta::nvertex, meta::on_dev);
    MEMCLEAN(odeg, meta::nvertex, meta::on_dev);

    MYLOG("degree data initialize complete");

    auto label_map = count_degree(src_raw, dst_raw, ideg, odeg, meta::nvertex, meta::nedge);

    MEMALLOC(IndexType, ideg_, label_map.size(), meta::on_dev);
    MEMALLOC(IndexType, odeg_, label_map.size(), meta::on_dev);

    auto r_label_map = prune_empty(meta::nvertex, label_map, ideg, odeg, 
                                   ideg_, odeg_);

    for ( int i=0; i<meta::nvertex; ++i ) {
      col_stride = std::max({col_stride, ideg[i]});
      row_stride = std::max({row_stride, odeg[i]});
    }

    MEMFREE(ideg, meta::on_dev);
    MEMFREE(odeg, meta::on_dev);
    ideg = ideg_; odeg = odeg_;
    meta::nvertex = label_map.size();
    
    MYLOG("degree counting complete, nvertex updated");

    _initializeField(CPU, meta::weighted, with_reverse);

    MYLOG("data initialize complete");

    for ( size_t i=0; i<meta::nedge; ++i )
    {
      IndexType s,d;
      s = r_label_map[src_raw[i]];
      d = r_label_map[dst_raw[i]];
      
      col_idx[col_stride * d + ideg[d]] = s;
      if ( with_weight ) col_eidx[col_stride * d + ideg[d]] = eidx_raw[i];
      ++ideg[d];
      if ( with_reverse )
      {
        row_idx[row_stride * s + odeg[s]] = d;
        if ( with_weight ) row_eidx[row_stride * s + odeg[s]] = eidx_raw[i];
        ++odeg[s];
      }
    }
  }

  ~ELLGraph() {
    if ( meta::independent ) {
      MEMFREE(col_idx, meta::on_dev);
      if ( meta::weighted ) { MEMFREE(col_idx, meta::on_dev); } 
      if ( row_idx != nullptr ) {
        MEMFREE(row_idx, meta::on_dev);
        if ( meta::weighted ) { MEMFREE(col_idx, meta::on_dev); }
      }
    }
  }

  // void toTensor(std::vector<torch::Tensor>& tlist) {
  //   ASSERT(meta::on_dev == CPU, "device copy is not ailable yet");
  //   ASSERT(tlist.size() == 4, "exact 6 input is required");
  //   meta::independent = false;
    
  //   torch::ScalarType indTp, valTp;
  //   if ( sizeof(IndexType) == 4 ) {
  //     indTp = torch::kI32;
  //   } else if ( sizeof(IndexType) == 8 ) {
  //     indTp = torch::kI64;
  //   } else { 
  //       ASSERT(false, "unsupported index type");
  //   }

  //   if (sizeof(IndexType) == 4) {
  //     valTp = torch::kF32;
  //   } else if (sizeof(IndexType) == 8) {
  //     valTp = torch::kF64;
  //   } else {
  //     ASSERT(false, "unsupported value type");
  //   }

  //   auto indOpt = torch::TensorOptions().dtype(indTp);
  //   auto valOpt = torch::TensorOptions().dtype(valTp);

  //   tlist[0] = torch::from_blob(col_idx, {meta::nvertex, col_stride}, indTp);
  //   if ( meta::weighted ) 
  //     tlist[1] = torch::from_blob(col_eidx, {meta::nvertex, col_stride}, valTp);
    
  //   if ( row_idx != nullptr ) {
  //     tlist[2] = torch::from_blob(row_idx, {meta::nvertex, row_stride}, indTp);
  //     if ( meta::weighted ) 
  //       tlist[3] = torch::from_blob(row_eidx, {meta::nvertex, col_stride}, valTp);
  //   }
  // }

  void cuda() {
    data_move(DeviceTag::CUDA);
  }

  void cpu() {
    data_move(DeviceTag::CPU);
  }
 
 protected:
  void _initializeField( DeviceTag dev = DeviceTag::CPU, bool with_weight = false, 
                         bool with_reverse = false ) {
    
    MEMALLOC(IndexType, col_idx, meta::nvertex*col_stride,dev);
    MEMSET(col_idx, meta::nvertex*col_stride, -1, dev);
    if ( with_weight ) {
      MEMALLOC(IndexType, col_eidx, meta::nvertex*col_stride, dev);
    } else {
      row_eidx = nullptr;
    }

    if ( with_reverse ) {
      MEMALLOC(IndexType, row_idx, meta::nvertex*row_stride, dev);
      MEMSET(row_idx, meta::nvertex*row_stride, -1, dev);
      if ( with_weight ) {
        MEMALLOC(IndexType, row_eidx, meta::nvertex*row_stride, dev);
      } else {
        row_eidx = nullptr;
      }
    } else {
      row_stride = 0;
      row_idx = nullptr;
      row_eidx = nullptr;
    }
  }

  void data_move( DeviceTag dst ) {
    if ( meta::on_dev == dst ) return;
    DeviceTag src = meta::on_dev;
    
    bool weighted = meta::weighted;
    bool reverse = ( row_idx != nullptr );
    size_t ell_size = meta::nvertex * col_stride;
    size_t rll_size = meta::nvertex * row_stride;

    IndexType *ci, *ri;
    IndexType *cd, *rd;

    MEMALLOC(IndexType, ci, ell_size, dst);
    MEMCOPY(ci, col_idx, ell_size, dst, src);
    MEMFREE(col_idx, src);
    col_idx = ci;

    if ( weighted ) { 
      MEMALLOC(IndexType, cd, ell_size, dst); 
      MEMCOPY(cd, col_eidx, ell_size, dst, src);
      MEMFREE(col_eidx, src);
      col_eidx = cd;
    }

    if ( reverse ) {
      MEMALLOC(IndexType, ri, rll_size, dst);
      MEMCOPY(ri, row_idx, rll_size, dst, src);
      MEMFREE(row_idx, src);
      row_idx = ri;
      if ( weighted ) { 
        MEMALLOC(IndexType, rd, rll_size, dst); 
        MEMCOPY(rd, row_eidx, rll_size, dst, src);
        MEMFREE(row_eidx, src);
        row_eidx = rd;
      }
    }
  }
};

}
#endif